from typing import Dict, Any, Optional, Tuple
from datetime import datetime, timedelta
from utils.logger import Logger
from utils.exceptions import RateLimitError
from middleware.redis_adapter import get_redis_adapter
from config import security_config
from utils.constants import DEFAULT_RATE_LIMIT

# 初始化日志
logger = Logger.get_logger(agent_name="distributed_rate_limiter")


class DistributedRateLimiter:
    """分布式限流（基于 Redis 滑动窗口算法）"""

    def __init__(self):
        self.redis_client = get_redis_adapter()  # 使用Redis适配器，支持集群模式
        self.prefix: str = "distributed:rate_limit:"  # 限流键前缀
        # 从配置读取限流规则（优先使用 security_config，其次用 constants 默认值）
        self.rate_limit_rules: Dict[str, Dict[str, int]] = self._load_rate_limit_rules()
        logger.info("分布式限流初始化完成，加载规则: {}".format(
            {k: f"{v['limit']}次/{v['period']}s" for k, v in self.rate_limit_rules.items()}
        ))

    def _load_rate_limit_rules(self) -> Dict[str, Dict[str, int]]:
        """加载限流规则（从 config 读取，兜底用 constants）"""
        config_rules = security_config.rate_limit.dict()
        return {
            "global": config_rules.get("global_limit", DEFAULT_RATE_LIMIT["api"]),
            "user": config_rules.get("user_limit", DEFAULT_RATE_LIMIT["user"]),
            "agent": config_rules.get("agent_limit", DEFAULT_RATE_LIMIT["agent"]),
            "hot_key": config_rules.get("hot_key_limit", {"limit": 1000, "period": 1})
        }

    def _get_key(self, limit_type: str, identifier: str = "") -> str:
        """生成限流 Redis 键名"""
        # 全局限流无需 identifier，其他维度需拼接唯一标识（如 user_id/agent_id）
        if limit_type == "global":
            return f"{self.prefix}global"
        return f"{self.prefix}{limit_type}:{identifier}"

    def _get_rule(self, limit_type: str) -> Tuple[int, int]:
        """获取限流规则（limit: 最大次数，period: 时间窗口秒）"""
        rule = self.rate_limit_rules.get(limit_type)
        if not rule:
            raise ValueError(f"无效的限流类型: {limit_type}，支持：{list(self.rate_limit_rules.keys())}")
        return rule["limit"], rule["period"]

    def _slide_window_check(self, key: str, limit: int, period: int) -> Tuple[bool, int, int]:
        """
        滑动窗口限流核心逻辑
        Args:
            key: 限流键
            limit: 最大请求次数
            period: 时间窗口（秒）
        Returns:
            (是否允许请求, 剩余次数, 重置时间秒)
        """
        now = datetime.now().timestamp()
        window_start = now - period  # 窗口起始时间（当前时间 - 时间窗口）
        pipeline = self.redis_client.pipeline()

        try:
            if pipeline:
                # 使用pipeline批量执行
                # 1. 删除窗口外的请求记录（ZSET 按分数排序，分数为时间戳）
                pipeline.zremrangebyscore(key, 0, window_start)
                # 2. 统计当前窗口内的请求次数
                pipeline.zcard(key)
                # 3. 添加当前请求记录（分数=当前时间戳，值=唯一标识避免重复）
                pipeline.zadd(key, {str(now): now})
                # 4. 设置 ZSET 过期时间（避免内存溢出）
                pipeline.expire(key, period + 1)  # 窗口+1秒确保所有记录都被删除
                # 执行管道操作
                results = pipeline.execute()
                current_count = results[1]  # 窗口内请求次数
            else:
                # 降级方案：单独执行每个命令
                # 1. 删除窗口外的请求记录
                self.redis_client.zremrangebyscore(key, 0, window_start)
                # 2. 统计当前窗口内的请求次数
                current_count = self.redis_client.zcard(key)
                # 3. 添加当前请求记录
                self.redis_client.zadd(key, {str(now): now})
                # 4. 设置过期时间
                self.redis_client.expire(key, period + 1)

            if current_count < limit:
                allowed = True
                remaining = limit - current_count - 1  # -1 是当前请求
            else:
                allowed = False
                remaining = 0

            # 计算重置时间（秒）
            reset_after = int(period - (now - window_start))
            return allowed, remaining, reset_after
        except Exception as e:
            logger.error(f"滑动窗口限流异常: key={key}, error={str(e)}", exc_info=True)
            # 限流检查失败时默认允许请求（避免影响业务）
            return True, -1, 0

    def check(self, limit_type: str, identifier: str = "") -> Dict[str, Any]:
        """
        检查限流（不抛出异常，返回结果）
        Args:
            limit_type: 限流类型（global/user/agent/hot_key）
            identifier: 唯一标识（user_id/agent_id/hot_key，全局限流无需传）
        Returns:
            {"allowed": bool, "remaining": int, "reset_after": int, "limit": int}
        """
        try:
            limit, period = self._get_rule(limit_type)
            key = self._get_key(limit_type, identifier)
            allowed, remaining, reset_after = self._slide_window_check(key, limit, period)

            result = {
                "allowed": allowed,
                "remaining": remaining,
                "reset_after": reset_after,
                "limit": limit,
                "period": period
            }

            if not allowed:
                logger.warning(
                    f"触发限流: type={limit_type}, id={identifier or 'global'}, "
                    f"count={limit - remaining}/{limit}, reset_after={reset_after}s"
                )
            else:
                logger.debug(
                    f"限流检查通过: type={limit_type}, id={identifier or 'global'}, "
                    f"remaining={remaining}, reset_after={reset_after}s"
                )
            return result
        except Exception as e:
            logger.error(f"限流检查异常: type={limit_type}, id={identifier}, error={str(e)}", exc_info=True)
            return {"allowed": True, "remaining": -1, "reset_after": 0, "limit": -1}

    def enforce(self, limit_type: str, identifier: str = "") -> None:
        """检查限流（触发限流时抛出异常）"""
        result = self.check(limit_type, identifier)
        if not result["allowed"]:
            raise RateLimitError(
                message=f"触发{limit_type}级限流，{result['limit']}次/{result['period']}s，剩余{result['remaining']}次",
                context={
                    "limit_type": limit_type,
                    "identifier": identifier,
                    "reset_after": result["reset_after"]
                }
            )

    def update_rule(self, limit_type: str, limit: int, period: int) -> bool:
        """动态更新限流规则（支持运行时调整）"""
        if limit_type not in self.rate_limit_rules:
            logger.warning(f"无效的限流类型: {limit_type}，无法更新规则")
            return False
        try:
            self.rate_limit_rules[limit_type] = {"limit": limit, "period": period}
            logger.info(f"更新限流规则: type={limit_type}, limit={limit}, period={period}s")
            return True
        except Exception as e:
            logger.error(f"更新限流规则异常: type={limit_type}, error={str(e)}", exc_info=True)
            return False


# 单例实例（全局复用）
rate_limiter = DistributedRateLimiter()

# 测试入口
if __name__ == "__main__":
    test_logger = Logger.update_context(task_id="rate_limiter_test")

    # 测试全局限流（100次/60秒）
    for i in range(5):
        try:
            rate_limiter.enforce("global")
            test_logger.info(f"全局限流第{i + 1}次请求: 允许")
        except RateLimitError as e:
            test_logger.error(f"全局限流第{i + 1}次请求: {e.message}")

    # 测试用户级限流（50次/60秒）
    user_id = "user_001"
    for i in range(5):
        try:
            rate_limiter.enforce("user", user_id)
            test_logger.info(f"用户{user_id}第{i + 1}次请求: 允许")
        except RateLimitError as e:
            test_logger.error(f"用户{user_id}第{i + 1}次请求: {e.message}")

    # 测试动态更新规则
    rate_limiter.update_rule("agent", limit=10, period=60)
    agent_id = "text_qa_001"
    for i in range(12):
        try:
            rate_limiter.enforce("agent", agent_id)
            test_logger.info(f"Agent{agent_id}第{i + 1}次请求: 允许")
        except RateLimitError as e:
            test_logger.error(f"Agent{agent_id}第{i + 1}次请求: {e.message}")